
package cn.cloud.all.security.oauth2.provider.endpoint;

import org.springframework.security.oauth2.common.exceptions.InvalidGrantException;
import org.springframework.security.oauth2.common.exceptions.InvalidRequestException;
import org.springframework.security.oauth2.common.exceptions.OAuth2Exception;
import org.springframework.security.oauth2.common.exceptions.RedirectMismatchException;
import org.springframework.security.oauth2.provider.ClientDetails;
import org.springframework.util.Assert;
import org.springframework.util.MultiValueMap;
import org.springframework.util.StringUtils;
import org.springframework.web.util.UriComponents;
import org.springframework.web.util.UriComponentsBuilder;

import java.util.*;

public class DefaultRedirectResolver implements RedirectResolver {

    private Collection<String> redirectGrantTypes = Arrays.asList("implicit", "authorization_code");

    private boolean matchSubdomains = false;

    private boolean matchPorts = true;

    public void setMatchSubdomains(boolean matchSubdomains) {
        this.matchSubdomains = matchSubdomains;
    }

    public void setMatchPorts(boolean matchPorts) {
        this.matchPorts = matchPorts;
    }

    public void setRedirectGrantTypes(Collection<String> redirectGrantTypes) {
        this.redirectGrantTypes = new HashSet<String>(redirectGrantTypes);
    }

    public String resolveRedirect(String requestedRedirect, ClientDetails client) throws OAuth2Exception {
        Set<String> authorizedGrantTypes = client.getAuthorizedGrantTypes();
        if (authorizedGrantTypes.isEmpty()) {
            throw new InvalidGrantException("A client must have at least one authorized grant type.");
        }
        if (!containsRedirectGrantType(authorizedGrantTypes)) {
            throw new InvalidGrantException("A redirect_uri can only be used by implicit or authorization_code grant types.");
        }

        Set<String> registeredRedirectUris = client.getRegisteredRedirectUri();
        if (registeredRedirectUris == null || registeredRedirectUris.isEmpty()) {
            throw new InvalidRequestException("At least one redirect_uri must be registered with the client.");
        }
        return obtainMatchingRedirect(registeredRedirectUris, requestedRedirect);
    }

    private boolean containsRedirectGrantType(Set<String> grantTypes) {
        for (String type : grantTypes) {
            if (redirectGrantTypes.contains(type)) {
                return true;
            }
        }
        return false;
    }

    protected boolean redirectMatches(String requestedRedirect, String redirectUri) {
        UriComponents requestedRedirectUri = UriComponentsBuilder.fromUriString(requestedRedirect).build();
        UriComponents registeredRedirectUri = UriComponentsBuilder.fromUriString(redirectUri).build();

        boolean schemeMatch = isEqual(registeredRedirectUri.getScheme(), requestedRedirectUri.getScheme());
        boolean userInfoMatch = isEqual(registeredRedirectUri.getUserInfo(), requestedRedirectUri.getUserInfo());
        boolean hostMatch = hostMatches(registeredRedirectUri.getHost(), requestedRedirectUri.getHost());
        boolean portMatch = !matchPorts || registeredRedirectUri.getPort() == requestedRedirectUri.getPort();
        boolean pathMatch = isEqual(registeredRedirectUri.getPath(), StringUtils.cleanPath(requestedRedirectUri.getPath()));
        boolean queryParamMatch = matchQueryParams(registeredRedirectUri.getQueryParams(), requestedRedirectUri.getQueryParams());

        return schemeMatch && userInfoMatch && hostMatch && portMatch && pathMatch && queryParamMatch;
    }

    private boolean matchQueryParams(MultiValueMap<String, String> registeredRedirectUriQueryParams, MultiValueMap<String, String> requestedRedirectUriQueryParams) {

        for (String key : registeredRedirectUriQueryParams.keySet()) {
            List<String> registeredRedirectUriQueryParamsValues = registeredRedirectUriQueryParams.get(key);
            List<String> requestedRedirectUriQueryParamsValues = requestedRedirectUriQueryParams.get(key);

            if (!registeredRedirectUriQueryParamsValues.equals(requestedRedirectUriQueryParamsValues)) {
                return false;
            }
        }

        return true;
    }

    private boolean isEqual(String str1, String str2) {
        return StringUtils.isEmpty(str1) && StringUtils.isEmpty(str2) || !StringUtils.isEmpty(str1) && str1.equals(str2);
    }

    private boolean hostMatches(String registered, String requested) {
        if (matchSubdomains) {
            return isEqual(registered, requested) || (requested != null && requested.endsWith("." + registered));
        }
        return isEqual(registered, requested);
    }

    private String obtainMatchingRedirect(Set<String> redirectUris, String requestedRedirect) {
        Assert.notEmpty(redirectUris, "Redirect URIs cannot be empty");

        if (redirectUris.size() == 1 && requestedRedirect == null) {
            return redirectUris.iterator().next();
        }

        for (String redirectUri : redirectUris) {
            if (requestedRedirect != null && redirectMatches(requestedRedirect, redirectUri)) {
                // Initialize with the registered redirect-uri
                UriComponentsBuilder redirectUriBuilder = UriComponentsBuilder.fromUriString(redirectUri);

                UriComponents requestedRedirectUri = UriComponentsBuilder.fromUriString(requestedRedirect).build();

                if (this.matchSubdomains) {
                    redirectUriBuilder.host(requestedRedirectUri.getHost());
                }
                if (!this.matchPorts) {
                    redirectUriBuilder.port(requestedRedirectUri.getPort());
                }
                redirectUriBuilder.replaceQuery(requestedRedirectUri.getQuery());        // retain additional params (if any)
                redirectUriBuilder.fragment(null);
                return redirectUriBuilder.build().toUriString();
            }
        }

        throw new RedirectMismatchException("Invalid redirect: " + requestedRedirect + " does not match one of the registered values.");
    }
}
